"""
FBCCA SSVEP分类器 - 适配102受试者数据集
从实时代码中提取的SSVEP信号处理和分类部分
"""

import numpy as np
import math
from scipy import signal
from scipy.signal import butter, filtfilt, cheby1, cheb1ord
from sklearn.cross_decomposition import CCA
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
from typing import List, Tuple, Optional
import time

class FBCCAClassifier:
    """FBCCA SSVEP分类器"""
    
    def __init__(self, fs: int = 250, num_harmonics: int = 3, num_fb: int = 5, 
                 window_length: float = 2.0, ssvep_freqs: List[float] = None):
        """
        初始化FBCCA分类器
        
        参数:
        fs: 采样频率 (Hz)
        num_harmonics: 谐波数量
        num_fb: 滤波器组数量  
        window_length: 分析窗口长度 (秒)
        ssvep_freqs: SSVEP频率列表
        """
        self.fs = fs
        self.num_harmonics = num_harmonics
        self.num_fb = num_fb
        self.window_length = window_length
        self.T = int(fs * window_length)  # 窗口采样点数
        
        # 默认12个SSVEP频率 (需要从实际数据集中确认)
        if ssvep_freqs is None:
            self.ssvep_freqs = [8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5, 12.0, 12.5, 13.0, 13.5]
        else:
            self.ssvep_freqs = ssvep_freqs
        
        self.n_classes = len(self.ssvep_freqs)
        
        print(f"FBCCA分类器初始化:")
        print(f"  采样率: {self.fs} Hz")
        print(f"  窗口长度: {self.window_length} s ({self.T} 采样点)")
        print(f"  谐波数: {self.num_harmonics}")
        print(f"  滤波器组数: {self.num_fb}")
        print(f"  SSVEP频率: {self.ssvep_freqs}")
        print(f"  类别数: {self.n_classes}")
    
    def get_reference_signals(self, target_freq: float) -> np.ndarray:
        """
        生成参考信号
        
        参数:
        target_freq: 目标频率
        
        返回:
        reference_signals: 参考信号矩阵 [n_harmonics*2, n_samples]
        """
        reference_signals = []
        t = np.arange(0, self.window_length, 1.0/self.fs)[:self.T]
        
        # 生成谐波的正弦和余弦分量
        for h in range(1, self.num_harmonics + 1):
            reference_signals.append(np.sin(2 * np.pi * h * target_freq * t))
            reference_signals.append(np.cos(2 * np.pi * h * target_freq * t))
        
        return np.array(reference_signals)
    
    def filter_bank(self, eeg_data: np.ndarray) -> np.ndarray:
        """
        滤波器组处理
        
        参数:
        eeg_data: EEG数据 [n_trials, n_channels, n_samples] 或 [n_channels, n_samples]
        
        返回:
        filtered_data: 滤波后的数据 [n_trials, n_fb, n_channels, n_samples]
        """
        # 处理输入维度
        if len(eeg_data.shape) == 2:
            eeg_data = eeg_data[np.newaxis, :, :]
        
        n_trials, n_channels, n_samples = eeg_data.shape
        filtered_data = np.zeros((n_trials, self.num_fb, n_channels, n_samples))
        
        nyq = self.fs / 2
        
        # 定义滤波器组参数 (适配SSVEP频率范围)
        passband = [6, 10, 14, 18, 22]  # 通带起始频率
        stopband = [4, 8, 12, 16, 20]   # 阻带起始频率
        highcut_pass, highcut_stop = 40, 50  # 高频截止
        
        gpass, gstop, Rp = 3, 40, 0.5
        
        for fb_idx in range(self.num_fb):
            try:
                # 设计Chebyshev I型带通滤波器
                Wp = [passband[fb_idx] / nyq, highcut_pass / nyq]
                Ws = [stopband[fb_idx] / nyq, highcut_stop / nyq]
                
                # 检查频率范围有效性
                if any(w >= 1.0 for w in Wp + Ws) or any(w <= 0 for w in Wp + Ws):
                    print(f"警告: 滤波器组 {fb_idx} 频率范围无效，跳过滤波")
                    filtered_data[:, fb_idx, :, :] = eeg_data
                    continue
                
                # 设计滤波器
                [N, Wn] = cheb1ord(Wp, Ws, gpass, gstop)
                [B, A] = cheby1(N, Rp, Wn, 'bandpass')
                
                # 应用滤波器
                for trial_idx in range(n_trials):
                    for ch_idx in range(n_channels):
                        try:
                            filtered_data[trial_idx, fb_idx, ch_idx, :] = filtfilt(
                                B, A, eeg_data[trial_idx, ch_idx, :],
                                padlen=min(3 * (max(len(B), len(A)) - 1), n_samples // 4)
                            )
                        except Exception as e:
                            print(f"滤波器组 {fb_idx} 通道 {ch_idx} 试验 {trial_idx} 滤波失败: {e}")
                            filtered_data[trial_idx, fb_idx, ch_idx, :] = eeg_data[trial_idx, ch_idx, :]
                        
            except Exception as e:
                print(f"滤波器组 {fb_idx} 设计失败: {e}")
                filtered_data[:, fb_idx, :, :] = eeg_data
        
        return filtered_data
    
    def find_correlation(self, X: np.ndarray, Y: np.ndarray) -> float:
        """
        计算CCA相关系数
        
        参数:
        X: EEG数据 [n_channels, n_samples]
        Y: 参考信号 [n_references, n_samples]
        
        返回:
        corr: 最大相关系数
        """
        try:
            cca = CCA(n_components=1)
            cca.fit(X.T, Y.T)
            x_a, y_b = cca.transform(X.T, Y.T)
            corr = np.corrcoef(x_a[:, 0], y_b[:, 0])[0, 1]
            
            # 处理NaN值
            if np.isnan(corr):
                corr = 0.0
                
            return corr
        except Exception as e:
            print(f"CCA计算失败: {e}")
            return 0.0
    
    def fbcca_classify(self, test_data: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        FBCCA分类
        
        参数:
        test_data: 测试数据 [n_trials, n_channels, n_samples] 或 [n_channels, n_samples]
        
        返回:
        predictions: 预测标签 [n_trials]
        scores: 分类得分 [n_trials]
        """
        # 处理输入维度
        single_trial = False
        if len(test_data.shape) == 2:
            test_data = test_data[np.newaxis, :, :]
            single_trial = True
        
        # 应用滤波器组
        filtered_data = self.filter_bank(test_data)
        n_trials = test_data.shape[0]
        
        # 滤波器组权重 (按频率递减)
        fb_weights = [math.pow(i, -1.25) + 0.25 for i in range(1, self.num_fb + 1)]
        
        predictions = []
        scores = []
        
        for trial_idx in range(n_trials):
            freq_scores = np.zeros(self.n_classes)
            
            # 对每个SSVEP频率计算得分
            for freq_idx, target_freq in enumerate(self.ssvep_freqs):
                # 生成参考信号
                reference = self.get_reference_signals(target_freq)
                
                # 计算滤波器组加权相关系数
                weighted_corr = 0.0
                for fb_idx in range(self.num_fb):
                    # 提取当前试验和滤波器组的数据
                    trial_data = filtered_data[trial_idx, fb_idx, :, :]
                    
                    # 计算CCA相关系数
                    corr = self.find_correlation(trial_data, reference)
                    
                    # 加权累加
                    weighted_corr += fb_weights[fb_idx] * (corr ** 2)
                
                freq_scores[freq_idx] = weighted_corr
            
            # 选择得分最高的频率
            predicted_freq_idx = np.argmax(freq_scores)
            max_score = np.max(freq_scores)
            
            predictions.append(predicted_freq_idx)
            scores.append(max_score)
        
        predictions = np.array(predictions)
        scores = np.array(scores)
        
        # 如果输入是单个试验，返回标量
        if single_trial:
            return predictions[0], scores[0]
        
        return predictions, scores
    
    def fit(self, X_train: np.ndarray, y_train: np.ndarray):
        """
        训练分类器 (FBCCA是无监督的，这里主要用于验证)
        
        参数:
        X_train: 训练数据 [n_trials, n_channels, n_samples]
        y_train: 训练标签 [n_trials]
        """
        print(f"FBCCA训练 (验证模式):")
        print(f"  训练数据: {X_train.shape}")
        print(f"  训练标签: {y_train.shape}")
        print(f"  标签分布: {np.bincount(y_train)}")
        
        # FBCCA是无监督的，不需要训练
        # 这里可以用训练数据验证分类器性能
        predictions, _ = self.fbcca_classify(X_train)
        accuracy = accuracy_score(y_train, predictions)
        print(f"  训练集准确率: {accuracy:.4f}")
        
        return self
    
    def predict(self, X_test: np.ndarray) -> np.ndarray:
        """
        预测标签
        
        参数:
        X_test: 测试数据 [n_trials, n_channels, n_samples]
        
        返回:
        predictions: 预测标签 [n_trials]
        """
        predictions, _ = self.fbcca_classify(X_test)
        return predictions
    
    def predict_proba(self, X_test: np.ndarray) -> np.ndarray:
        """
        预测概率分数
        
        参数:
        X_test: 测试数据 [n_trials, n_channels, n_samples]
        
        返回:
        scores: 分类得分 [n_trials]
        """
        _, scores = self.fbcca_classify(X_test)
        return scores


class SSVEPPreprocessor:
    """SSVEP数据预处理器"""
    
    def __init__(self, fs: int = 250, low_freq: float = 3.0, high_freq: float = 50.0):
        """
        初始化预处理器
        
        参数:
        fs: 采样频率
        low_freq: 低频截止
        high_freq: 高频截止
        """
        self.fs = fs
        self.low_freq = low_freq
        self.high_freq = high_freq
        
        print(f"SSVEP预处理器初始化:")
        print(f"  采样率: {fs} Hz")
        print(f"  带通滤波: {low_freq}-{high_freq} Hz")
    
    def bandpass_filter(self, data: np.ndarray) -> np.ndarray:
        """
        带通滤波
        
        参数:
        data: 输入数据 [n_trials, n_channels, n_samples] 或 [n_channels, n_samples]
        
        返回:
        filtered_data: 滤波后的数据
        """
        original_shape = data.shape
        
        # 处理输入维度
        if len(data.shape) == 2:
            data = data[np.newaxis, :, :]
            single_trial = True
        else:
            single_trial = False
        
        n_trials, n_channels, n_samples = data.shape
        filtered_data = np.zeros_like(data)
        
        # 设计Butterworth带通滤波器
        nyq = self.fs / 2
        low_norm = self.low_freq / nyq
        high_norm = self.high_freq / nyq
        
        # 检查频率范围
        if low_norm <= 0 or high_norm >= 1 or low_norm >= high_norm:
            print(f"警告: 滤波器频率范围无效，跳过滤波")
            return data.reshape(original_shape)
        
        try:
            b, a = butter(4, [low_norm, high_norm], btype='band')
            
            for trial_idx in range(n_trials):
                for ch_idx in range(n_channels):
                    if n_samples >= 3 * max(len(a), len(b)):
                        filtered_data[trial_idx, ch_idx, :] = filtfilt(b, a, data[trial_idx, ch_idx, :])
                    else:
                        print(f"警告: 试验 {trial_idx} 通道 {ch_idx} 数据点不足，跳过滤波")
                        filtered_data[trial_idx, ch_idx, :] = data[trial_idx, ch_idx, :]
            
            print(f"带通滤波完成: {self.low_freq}-{self.high_freq} Hz")
            
        except Exception as e:
            print(f"带通滤波失败: {e}")
            filtered_data = data
        
        return filtered_data.reshape(original_shape)
    
    def baseline_correction(self, data: np.ndarray, baseline_samples: int = 125) -> np.ndarray:
        """
        基线校正
        
        参数:
        data: 输入数据 [n_trials, n_channels, n_samples]
        baseline_samples: 基线样本点数
        
        返回:
        corrected_data: 基线校正后的数据
        """
        if len(data.shape) != 3:
            raise ValueError("基线校正需要3维数据 [n_trials, n_channels, n_samples]")
        
        # 计算基线 (前baseline_samples个样本点的平均值)
        baseline = np.mean(data[:, :, :baseline_samples], axis=2, keepdims=True)
        
        # 基线校正
        corrected_data = data - baseline
        
        print(f"基线校正完成，使用前 {baseline_samples} 个样本点")
        return corrected_data
    
    def preprocess(self, data: np.ndarray, apply_bandpass: bool = True, 
                  apply_baseline: bool = True, baseline_samples: int = 125) -> np.ndarray:
        """
        完整预处理流程
        
        参数:
        data: 输入数据
        apply_bandpass: 是否应用带通滤波
        apply_baseline: 是否应用基线校正
        baseline_samples: 基线样本点数
        
        返回:
        processed_data: 预处理后的数据
        """
        processed_data = data.copy()
        
        if apply_bandpass:
            processed_data = self.bandpass_filter(processed_data)
        
        if apply_baseline and len(processed_data.shape) == 3:
            processed_data = self.baseline_correction(processed_data, baseline_samples)
        
        return processed_data


def evaluate_classifier(classifier, X_test: np.ndarray, y_test: np.ndarray, 
                       class_names: List[str] = None) -> dict:
    """
    评估分类器性能
    
    参数:
    classifier: 训练好的分类器
    X_test: 测试数据
    y_test: 测试标签
    class_names: 类别名称列表
    
    返回:
    results: 评估结果字典
    """
    print("开始评估分类器性能...")
    
    # 预测
    start_time = time.time()
    predictions = classifier.predict(X_test)
    prediction_time = time.time() - start_time
    
    # 计算准确率
    accuracy = accuracy_score(y_test, predictions)
    
    # 计算混淆矩阵
    cm = confusion_matrix(y_test, predictions)
    
    # 分类报告
    if class_names is None:
        class_names = [f"Class_{i}" for i in range(len(classifier.ssvep_freqs))]
    
    report = classification_report(y_test, predictions, target_names=class_names, output_dict=True, zero_division=0)
    
    # 计算ITR (Information Transfer Rate)
    n_classes = len(classifier.ssvep_freqs)
    if accuracy > 1/n_classes:  # 避免log(0)
        itr = (np.log2(n_classes) + accuracy * np.log2(accuracy) + 
               (1 - accuracy) * np.log2((1 - accuracy) / (n_classes - 1))) * 60  # bits/min
    else:
        itr = 0.0
    
    results = {
        'accuracy': accuracy,
        'confusion_matrix': cm,
        'classification_report': report,
        'itr': itr,
        'prediction_time': prediction_time,
        'n_test_samples': len(y_test)
    }
    
    print(f"评估结果:")
    print(f"  准确率: {accuracy:.4f}")
    print(f"  ITR: {itr:.2f} bits/min")
    print(f"  预测时间: {prediction_time:.4f}s ({len(y_test)} 个样本)")
    print(f"  平均预测时间: {prediction_time/len(y_test)*1000:.2f}ms/样本")
    
    return results


def main():
    """测试FBCCA分类器"""
    # 创建模拟数据进行测试
    fs = 250
    window_length = 2.0
    n_channels = 8
    n_samples = int(fs * window_length)
    n_trials = 120  # 10个区块 * 12个目标
    
    # 创建分类器
    ssvep_freqs = [8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5, 12.0, 12.5, 13.0, 13.5]
    classifier = FBCCAClassifier(fs=fs, ssvep_freqs=ssvep_freqs, window_length=window_length)
    
    # 创建预处理器
    preprocessor = SSVEPPreprocessor(fs=fs)
    
    print("="*50)
    print("FBCCA分类器测试完成")
    print(f"支持 {len(ssvep_freqs)} 个SSVEP频率")
    print(f"窗口长度: {window_length}s ({n_samples} 采样点)")


if __name__ == "__main__":
    main()
